import base64
import io
import os
import sys

import gradio as gr

from sdk.httpclient import CTClientBuilder, CTClient

base_url = "https://ai-global.ctapi.ctyun.cn"

client: CTClient


def _get_image_base64(img):
    """
    返回图片的base64大小（单位MB）和内容
    """
    byte_arr = io.BytesIO()
    img.save(byte_arr, format='PNG')
    byte_arr = byte_arr.getvalue()
    base64_str = base64.b64encode(byte_arr).decode('utf-8')
    base64_mb_size = sys.getsizeof(base64_str) / (1024 * 1024)
    return base64_mb_size, base64_str


def image_porn(img):
    base64_size, base64_str = _get_image_base64(img)
    if base64_size > 2.0:
        return {"errMsg": f"上传的图片进行base64编码后大小为 {base64_size:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2f820wzpai9s/PictureIdentity/api/v1/image_porn.json"
    response = client.request("POST", url, body={"data": [base64_str]})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return result
    except Exception as e:
        print(f"error: {e}")
        return None, response_json


def image_violence(img):
    base64_size, base64_str = _get_image_base64(img)
    if base64_size > 2.0:
        return {"errMsg": f"上传的图片进行base64编码后大小为 {base64_size:.2f} MB，超过 2.00 MB"}, None

    url = f"{base_url}/v1/aiop/api/2f3p4ejokcg0/ViolencePhoto/api/v1/image_violence.json"
    response = client.request("POST", url, body={"data": [base64_str]})
    response.raise_for_status()
    response_json = response.json()
    try:
        result = response_json['returnObj']
        return result
    except Exception as e:
        print(f"error: {e}")
        return None, response_json


def text_porn(text):
    url = f"{base_url}/v1/aiop/api/2f5o7mk00yrk/abc/api/v1/text_porn.json"
    response = client.request("POST", url, body={"data": [text]})
    response.raise_for_status()
    response_json = response.json()
    try:
        return response_json['returnObj']
    except:
        return response_json


def text_politic(text):
    url = f"{base_url}/v1/aiop/api/2f3rrma38pvk/FileIdentity2/api/v1/text_politic.json"
    response = client.request("POST", url, body={"data": [text]})
    response.raise_for_status()
    response_json = response.json()
    try:
        return response_json['returnObj']
    except:
        return response_json


product_intro = "进一步了解天翼云内容审核产品：<a href='https://www.ctyun.cn/products/contentsafe'>https://www.ctyun.cn/products/contentsafe</a>"
gr_config = {
    "image_porn": gr.Interface(fn=image_porn,
                                inputs=gr.Image(type="pil", label="上传图像"),
                                outputs=[
                                    gr.Textbox(label="鉴定结果")
                                ],
                                examples=[],
                                title="内容审核 - 图片鉴黄",
                                description="用于自动审核指定图片中是否存在色情内容，并给出该图片的涉黄指数，表明图片内容的涉黄程度。",
                                article=product_intro),
    "image_violence": gr.Interface(fn=image_violence,
                                inputs=gr.Image(type="pil", label="上传图像"),
                                outputs=[
                                    gr.Textbox(label="识别结果")
                                ],
                                examples=[],
                                title="内容审核 - 图片涉暴恐识别",
                                description="用于自动审核指定图片中是否存在涉暴涉恐内容，并给出该图片的涉暴涉恐指数，表明图片内容的涉暴涉恐程度。",
                                article=product_intro),
    "text_porn": gr.Interface(fn=text_porn,
                                inputs='text',
                                outputs=[
                                    gr.Textbox(label="审核结果")
                                ],
                                examples=[],
                                title="内容审核 - 文本涉黄审核",
                                description="用于自动审核指定文本中是否存在涉黄和色情内容，并给出该文本的涉黄指数，表明文本内容的涉黄程度。",
                                article=product_intro),
    "text_politic": gr.Interface(fn=text_politic,
                              inputs='text',
                              outputs=[
                                  gr.Textbox(label="审核结果")
                              ],
                              examples=[],
                              title="内容审核 - 文本涉政审核",
                              description="用于自动审核指定文本中是否存在涉政敏感内容，并给出该文本的涉政指数，表明文本内容的涉政程度。",
                              article=product_intro),
}


def get_not_empty_env(key):
    value = os.getenv(key, "").strip()
    if value == "":
        raise Exception(f"env {key} is not set or empty")
    return value


def run():
    ctyun_ak = get_not_empty_env("ext_cf_ctyun_ak")
    ctyun_sk = get_not_empty_env("ext_cf_ctyun_sk")
    ctyun_ai_app_key = get_not_empty_env("ext_cf_ctyun_ai_app_key")
    app_type = get_not_empty_env("ext_cf_app_type")

    global client
    client = CTClientBuilder().with_ak(ctyun_ak).with_sk(ctyun_sk).with_ai_app_key(ctyun_ai_app_key).build()

    gr_config[app_type].launch(server_name="0.0.0.0", server_port=9000)


if __name__ == '__main__':
    run()
